require 'base64'
require 'cgi'
require 'digest/sha1'
require 'net/http'
require 'openssl'

# Ruby class for verifying YubiKey OTPs using HMAC signature verification on
# the request and the response.

# Sample uses:
#   yubi = YubiSigned.new(otp)
#   puts yubi.unique_id                  #-> 'ccccccccdefg'
#   puts yubi.verify                     #-> 'OK'
#   puts yubi.verify                     #-> 'REPLAYED_OTP'
#   puts yubi.verify(new_otp)            #-> 'OK'
# Alternative quick verification:
#   puts YubiSigned.new.verify(otp)      #-> 'OK'

class YubiSigned

  # Client information. Change the next two lines to your own ID and secret
  CLIENT_ID = 1234
  SECRET_BASE64 = 'aBcDeFgHiJkLmNoPqRsTuVwXyZ'
  SECRET = Base64.decode64(SECRET_BASE64)

  # Address of Yubico WSAPI
  HOST = 'api.yubico.com'
  PATH = '/wsapi/verify'

  # verify() can return these codes in addition to Yubico status codes
  ERROR_BAD_OTP = 'BAD_OTP'
  ERROR_HTTP_FAILURE = 'HTTP_FAILURE'
  ERROR_BAD_YUBI_SIGNATURE = 'BAD_YUBI_SIGNATURE'
  ERROR_NO_STATUS = 'NO_STATUS'

  # Field for the signatures in query string and in response body
  SIGNATURE_FIELD = 'h'

  # The one-time password
  attr_reader :otp

  def initialize(otp = nil)
    set_otp(otp)
  end

  # The verify method returns 'OK' or an error string
  def verify(otp = nil)
    set_otp(otp)
    get_verification
  end

  # Unique ID of the device, = the first 12 characters of the OTP,
  def unique_id
    has_otp? ? @otp[0,12] : nil
  end
  
  private #-----

  # Save the OTP if it seems reasonable
  def set_otp(otp)
    @otp = otp.strip if otp.is_a?(String)
  end

  # Is there a reasonable OTP?
  def has_otp?
    @otp && @otp.length && @otp.length == 44
  end

  # Sign and send WSAPI request, check response signature, return the status
  def get_verification
    return ERROR_BAD_OTP unless has_otp?
    query_hash = {'id' => CLIENT_ID, 'otp' => @otp}
    path_q = PATH + "?" + hash_to_query_signed(query_hash)
    response = Net::HTTP.start(HOST) { |http| http.get path_q } rescue nil
    return ERROR_HTTP_FAILURE unless response.class == Net::HTTPOK
    body = response.body rescue ''
    return ERROR_BAD_YUBI_SIGNATURE unless signature_ok?(body)
    /status=(\S+)/ =~ body ? $1 : ERROR_NO_STATUS
  end

  # In: query string e.g. 'a=123&b=456'.
  # Out: HMAC SHA1 digest as string (base64 encoded)
  def make_signature(query)
    signature = OpenSSL::HMAC.digest(OpenSSL::Digest::SHA1.new, SECRET, query)
    Base64.encode64(signature).chomp
  end

  # In: a hash (keys must be strings) e.g. {'a' => 123, 'b' => 456}.
  # Out: a query string 'a=123&b=456'.
  def hash_to_query(hash)
    hash.sort.map{|x| "#{x[0].to_s}=#{x[1].to_s}"}.join('&')
  end

  # In: a hash (keys must be strings).
  # Out: a query string with signature.
  def hash_to_query_signed(hash)
    query = hash_to_query(hash)
    signature = make_signature(query)
    escaped_signature = CGI.escape(signature)
    "#{query}&#{SIGNATURE_FIELD}=#{escaped_signature}"
  end

  # In: the body of the HTTP response from Yubico.
  # Out: data as a hash (keys and values are still strings).
  def body_to_hash(text)
    text.split("\n").inject({}) do |memo, item|
      matches = item.chomp.match(/(.*?)=(.*)/)
      memo[matches[1]] = matches[2] if matches
      memo
    end
  end

  # In: the body of the HTTP response from Yubico.
  # Out: boolean: is the included signature valid for the rest of the data?
  def signature_ok?(body)
    body_hash = body_to_hash(body)
    signature_provided = body_hash.delete(SIGNATURE_FIELD)
    query = hash_to_query(body_hash)
    signature_provided == make_signature(query)
  end

end
